package com.cl.spark.node;

import com.alibaba.fastjson.JSONObject;
import com.cl.spark.base.BaseSparkNode;
import com.cl.spark.dto.SparkParam;
import com.cl.spark.dto.SparkResult;
import com.cl.spark.enums.SparkNodeEnum;
import com.cl.spark.util.StringUtil;
import org.apache.spark.sql.Dataset;
import org.apache.spark.sql.Row;
import org.apache.spark.sql.SparkSession;
import org.apache.spark.sql.functions;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;


@Component
public class SqlQueryNode extends BaseSparkNode {

    @Autowired
    SparkSession sparkSession;

    @Override
    public SparkResult process(SparkParam sparkParam) {
        JSONObject expression = sparkParam.getNodeExpression();
        String sql = String.valueOf(expression.get("sql"));
        String dbCode = String.valueOf(expression.get("dbCode"));

        String url = expression.getString("url");
        String username = expression.getString("username");
        String password = expression.getString("password");


        Dataset<Row> dataset = sparkSession.read()
                .format("jdbc")
                .option("url", url)
                .option("dbtable", "( " + sql + " ) as temp_" + System.currentTimeMillis())
                .option("user", username)
                .option("password", password)
                .load();
        dataset = dataset.withColumn("DB_CODE", functions.lit(dbCode)).withColumn("TABLE_NAME",functions.lit(StringUtil.extractSQLTableNames(sql)));

        return SparkResult.success(dataset);

    }

    @Override
    public SparkNodeEnum getType() {
        return SparkNodeEnum.SQL_QUERY;
    }


}
